import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from vmdpy import VMD

# Load real futures data
df = pd.read_csv('data/SF0SH.csv')
close_prices = df['close'].values
dates = pd.to_datetime(df['date_time'])
t = np.arange(len(close_prices))
signal = close_prices.astype(float)

# VMD parameters
alpha = 2000       # Moderate bandwidth constraint
tau = 0.            # Noise-tolerance (no noise here)
K = 3              # Number of modes
DC = 0             # No DC part in the signal
init_f = None      # Initialize center frequencies of modes
tol = 1e-7         # Tolerance for convergence

# Run VMD
u, u_hat, omega = VMD(signal, alpha, tau, K, DC, init_f, tol)

# Plot the original signal and the decomposed modes
plt.figure(figsize=(12, 8))
plt.subplot(K + 1, 1, 1)
plt.plot(dates, signal)
plt.title('SF0SH Futures Price')
plt.ylabel('Price')

for i in range(K):
    plt.subplot(K + 1, 1, i + 2)
    plt.plot(dates, u[i, :])
    plt.title(f'Decomposed Mode {i+1}')
    plt.ylabel('Price')

plt.xlabel('Date')
plt.tight_layout()
plt.show()
